# Copyright 2025 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Contains templates for Suzuki-Trotter approximation based subroutines.
"""

import numpy as np

from pennylane.estimator.compact_hamiltonian import (
    CDFHamiltonian,
    THCHamiltonian,
    VibrationalHamiltonian,
    VibronicHamiltonian,
)
from pennylane.estimator.ops.op_math.symbolic import Controlled, Prod
from pennylane.estimator.ops.qubit.non_parametric_ops import Hadamard, T, X
from pennylane.estimator.ops.qubit.parametric_ops_multi_qubit import MultiRZ
from pennylane.estimator.ops.qubit.parametric_ops_single_qubit import RZ
from pennylane.estimator.resource_operator import (
    CompressedResourceOp,
    GateCount,
    ResourceOperator,
    _dequeue,
    resource_rep,
)
from pennylane.estimator.wires_manager import Allocate, Deallocate
from pennylane.wires import Wires, WiresLike

from .subroutines import (
    BasisRotation,
    OutMultiplier,
    OutOfPlaceSquare,
    PhaseGradient,
    Select,
    SemiAdder,
)

# pylint: disable=arguments-differ, too-many-arguments, super-init-not-called


class TrotterProduct(ResourceOperator):
    r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
    exponential of a Hamiltonian operator.

    The Suzuki-Trotter product formula provides a method to approximate the matrix exponential of
    Hamiltonian expressed as a linear combination of terms which in general do not commute.
    Consider the Hamiltonian :math:`H = \Sigma^{N}_{j=0} O_{j}`: the product formula is constructed using
    symmetrized products of the terms in the Hamiltonian. The symmetrized products of order
    :math:`m \in [1, 2, 4, ..., 2k]` with :math:`k \in \mathbb{N}` are given by:

    .. math::

        \begin{align}
            S_{1}(t) &= \Pi_{j=0}^{N} \ e^{i t O_{j}} \\
            S_{2}(t) &= \Pi_{j=0}^{N} \ e^{i \frac{t}{2} O_{j}} \cdot \Pi_{j=N}^{0} \ e^{i \frac{t}{2} O_{j}} \\
            &\vdots \\
            S_{m}(t) &= S_{m-2}(p_{m}t)^{2} \cdot S_{m-2}((1-4p_{m})t) \cdot S_{m-2}(p_{m}t)^{2},
        \end{align}

    where the coefficient is :math:`p_{m} = 1 / (4 - \sqrt[m - 1]{4})`. The :math:`m^{\text{th}}` order,
    :math:`n`-step Suzuki-Trotter approximation is then defined as:

    .. math:: e^{iHt} \approx \left [S_{m}(t / n)  \right ]^{n}.

    For more details, see `J. Math. Phys. 32, 400 (1991) <https://pubs.aip.org/aip/jmp/article-abstract/32/2/400/229229>`_.

    Args:
        first_order_expansion (list[~pennylane.estimator.ResourceOperator]): A list of operators
            constituting the first order expansion of the Hamiltonian to be approximately exponentiated.
        num_steps (int): number of Trotter steps to perform
        order (int): order of the Suzuki-Trotter approximation; must be ``1`` or an even number
        wires (list[int] | None): The wires on which the operator acts. If provided, these wire
            labels will be used instead of the wires provided by the ResourceOperators in the
            :code:`first_order_expansion`.

    Resources:
        The resources are defined according to the recursive formula presented above.
        The number of times an operator :math:`e^{itO_{j}}` is applied depends on the
        number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

        .. math:: C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}

        Furthermore, because of the symmetric form of the recursive formula, the first and last terms are grouped.
        This reduces the counts for those terms to:

        .. math::

            \begin{align}
                C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
            \end{align}

    .. seealso:: The corresponding PennyLane operation :class:`~.TrotterProduct`

    .. seealso::
        :class:`~.estimator.templates.TrotterCDF`,
        :class:`~.estimator.templates.TrotterTHC`,
        :class:`~.estimator.templates.TrotterVibrational`,
        :class:`~.estimator.templates.TrotterVibronic`

    **Example**

    The resources for this operation are computed using:

    >>> import pennylane.estimator as qre
    >>> num_steps, order = (1, 2)
    >>> first_order_expansion = [qre.RX(), qre.RY()] # H = X + Y
    >>> gate_set = {"RX", "RY"}
    >>> res = qre.estimate(qre.TrotterProduct(first_order_expansion, num_steps, order), gate_set=gate_set)
    >>> print(res)
    --- Resources: ---
     Total wires: 1
        algorithmic wires: 1
        allocated wires: 0
             zero state: 0
             any state: 0
     Total gates : 3
      'RX': 2,
      'RY': 1
    """

    resource_keys = {"first_order_expansion", "num_steps", "order", "num_wires"}

    def __init__(
        self,
        first_order_expansion: list,
        num_steps: int,
        order: int,
        wires: WiresLike | None = None,
    ):

        _dequeue(op_to_remove=first_order_expansion)
        self.queue()

        try:
            cmpr_ops = tuple(op.resource_rep_from_op() for op in first_order_expansion)
        except AttributeError as error:
            raise ValueError(
                "All components of first_order_expansion must be instances of `ResourceOperator` in order to obtain resources."
            ) from error

        self.first_order_expansion = cmpr_ops
        self.num_steps = num_steps
        self.order = order

        if wires is not None:  # User defined wires take precedent
            self.wires = Wires(wires)
            self.num_wires = len(self.wires)

        else:  # Otherwise determine the wires from the ops in the first order expansion
            ops_wires = Wires.all_wires(
                [op.wires for op in first_order_expansion if op.wires is not None]
            )
            fewest_unique_wires = max(op.num_wires for op in cmpr_ops)

            if len(ops_wires) < fewest_unique_wires:  # If the expansion didn't provide enough wire
                self.wires = None  # labels we assume they all act on the same set
                self.num_wires = fewest_unique_wires

            else:  # If there are more wire labels, use that as the operator wires
                self.wires = ops_wires
                self.num_wires = len(self.wires)

    @property
    def resource_params(self) -> dict:
        r"""Returns a dictionary containing the minimal information needed to compute the resources.

        Returns:
            dict: A dictionary containing the resource parameters:
                * first_order_expansion (list[:class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`]): A list of operators,
                  in the compressed representation, constituting the first order expansion of the Hamiltonian to be approximately exponentiated.
                * num_steps (int): number of Trotter steps to perform
                * order (int): order of the Suzuki-Trotter approximation, must be 1 or even
                * num_wires (int): number of wires the operator acts on
        """
        return {
            "first_order_expansion": self.first_order_expansion,
            "num_steps": self.num_steps,
            "order": self.order,
            "num_wires": self.num_wires,
        }

    @classmethod
    def resource_rep(
        cls,
        first_order_expansion: list,
        num_steps: int,
        order: int,
        num_wires: int,
    ) -> CompressedResourceOp:
        """Returns a compressed representation containing only the parameters of
        the Operator that are needed to compute a resource estimation.

        Args:
            first_order_expansion (list[:class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`]): A list of operators,
                in the compressed representation, constituting
                the first order expansion of the Hamiltonian to be approximately exponentiated.
            num_steps (int): number of Trotter steps to perform
            order (int): order of the Suzuki-Trotter approximation, must be 1 or even
            num_wires (int): number of wires the operator acts on

        Returns:
            :class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`: the operator in a compressed representation
        """
        params = {
            "first_order_expansion": first_order_expansion,
            "num_steps": num_steps,
            "order": order,
            "num_wires": num_wires,
        }
        return CompressedResourceOp(cls, num_wires, params)

    @classmethod
    def resource_decomp(
        cls,
        first_order_expansion: list,
        num_steps: int,
        order: int,
        num_wires: int,  # pylint: disable=unused-argument
    ) -> list[GateCount]:
        r"""Returns a list representing the resources of the operator. Each object represents a
        quantum gate and the number of times it occurs in the decomposition.

        Args:
            first_order_expansion (list[:class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`]): A list of operators,
                in the compressed representation, constituting
                the first order expansion of the Hamiltonian to be approximately exponentiated.
            num_steps (int): number of Trotter steps to perform
            order (int): order of the Suzuki-Trotter approximation, must be 1 or even

        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.
        """
        k = order // 2
        gate_list = []

        if order == 1:
            for op in first_order_expansion:
                gate_list.append(GateCount(op, num_steps))
            return gate_list

        # For first and last fragment
        first_frag = first_order_expansion[0]
        last_frag = first_order_expansion[-1]
        gate_list.append(GateCount(first_frag, num_steps * (5 ** (k - 1)) + 1))
        gate_list.append(GateCount(last_frag, num_steps * (5 ** (k - 1))))

        # For rest of the fragments
        for op in first_order_expansion[1:-1]:
            gate_list.append(GateCount(op, 2 * num_steps * (5 ** (k - 1))))

        return gate_list


class TrotterCDF(ResourceOperator):
    r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
    exponential of a compressed double-factorized (CDF) Hamiltonian.

    The Suzuki-Trotter product formula provides a method to approximate the matrix exponential of
    Hamiltonian expressed as a linear combination of terms which in general do not commute.
    Consider the Hamiltonian :math:`H = \Sigma^{N}_{j=0} O_{j}`: the product formula is constructed using
    symmetrized products of the terms in the Hamiltonian. The symmetrized products of order
    :math:`m \in [1, 2, 4, ..., 2k]` with :math:`k \in \mathbb{N}` are given by:

    .. math::

        \begin{align}
            S_{1}(t) &= \Pi_{j=0}^{N} \ e^{i t O_{j}} \\
            S_{2}(t) &= \Pi_{j=0}^{N} \ e^{i \frac{t}{2} O_{j}} \cdot \Pi_{j=N}^{0} \ e^{i \frac{t}{2} O_{j}} \\
            &\vdots \\
            S_{m}(t) &= S_{m-2}(p_{m}t)^{2} \cdot S_{m-2}((1-4p_{m})t) \cdot S_{m-2}(p_{m}t)^{2},
        \end{align}

    where the coefficient is :math:`p_{m} = 1 / (4 - \sqrt[m - 1]{4})`. The :math:`m^{\text{th}}`
    order, :math:`n`-step Suzuki-Trotter approximation is then defined as:

    .. math::

        e^{iHt} \approx \left [S_{m}(t / n)  \right ]^{n}.

    For more details see `J. Math. Phys. 32, 400 (1991) <https://pubs.aip.org/aip/jmp/article-abstract/32/2/400/229229>`_.

    Args:
        cdf_ham (:class:`~.pennylane.estimator.compact_hamiltonian.CDFHamiltonian`):
            a compressed double factorized Hamiltonian to be approximately exponentiated
        num_steps (int): number of Trotter steps to perform
        order (int): order of the approximation, must be ``1`` or an even number
        wires (list[int] | None): the wires on which the operator acts

    Resources:
        The resources are defined according to the recursive formula presented above.
        The number of times an operator :math:`e^{itO_{j}}` is applied depends on the
        number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

        .. math::

            C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

        Furthermore, because of the symmetric form of the recursive formula, the first and last terms get grouped.
        This reduces the counts for those terms to:

        .. math::

            \begin{align}
                C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
            \end{align}

        The resources for a single step expansion of compressed double factorized Hamiltonian are
        calculated based on `arXiv:2506.15784 <https://arxiv.org/abs/2506.15784>`_.

    .. seealso::
        :class:`~.estimator.compact_hamiltonian.CDFHamiltonian`

    .. seealso:: :class:`~.TrotterProduct`

    **Example**

    The resources for this operation are computed using:

    >>> import pennylane.estimator as qre
    >>> num_steps, order = (1, 2)
    >>> cdf_ham = qre.CDFHamiltonian(num_orbitals = 4, num_fragments = 4)
    >>> res = qre.estimate(qre.TrotterCDF(cdf_ham, num_steps, order))
    >>> print(res)
    --- Resources: ---
     Total wires: 8
        algorithmic wires: 8
        allocated wires: 0
             zero state: 0
             any state: 0
     Total gates : 2.238E+4
      'T': 2.075E+4,
      'CNOT': 448,
      'Z': 336,
      'S': 504,
      'Hadamard': 336
    """

    resource_keys = {"cdf_ham", "num_steps", "order"}

    def __init__(
        self,
        cdf_ham: CDFHamiltonian,
        num_steps: int,
        order: int,
        wires: WiresLike | None = None,
    ):

        if not isinstance(cdf_ham, CDFHamiltonian):
            raise TypeError(
                f"Unsupported Hamiltonian representation for TrotterCDF."
                f"This method works with cdf Hamiltonian, {type(cdf_ham)} provided"
            )
        self.num_steps = num_steps
        self.order = order
        self.cdf_ham = cdf_ham

        self.num_wires = 2 * cdf_ham.num_orbitals

        if wires is not None and len(Wires(wires)) != self.num_wires:
            raise ValueError(f"Expected {self.num_wires} wires, got {len(Wires(wires))}")

        super().__init__(wires=wires)

    @property
    def resource_params(self) -> dict:
        r"""Returns a dictionary containing the minimal information needed to compute the resources.

        Returns:
            dict: A dictionary containing the resource parameters:
                * cdf_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.CDFHamiltonian`): a compressed double factorized
                  Hamiltonian to be approximately exponentiated
                * num_steps (int): number of Trotter steps to perform
                * order (int): order of the approximation, must be 1 or even.
        """
        return {
            "cdf_ham": self.cdf_ham,
            "num_steps": self.num_steps,
            "order": self.order,
        }

    @classmethod
    def resource_rep(
        cls, cdf_ham: CDFHamiltonian, num_steps: int, order: int
    ) -> CompressedResourceOp:
        """Returns a compressed representation containing only the parameters of
        the Operator that are needed to compute a resource estimation.

        Args:
            cdf_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.CDFHamiltonian`):
                a compressed double factorized Hamiltonian to be approximately exponentiated
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even.

        Returns:
            :class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`: the operator in a compressed representation
        """
        params = {
            "cdf_ham": cdf_ham,
            "num_steps": num_steps,
            "order": order,
        }
        num_wires = 2 * cdf_ham.num_orbitals
        return CompressedResourceOp(cls, num_wires, params)

    @classmethod
    def resource_decomp(
        cls, cdf_ham: CDFHamiltonian, num_steps: int, order: int
    ) -> list[GateCount]:
        r"""Returns a list representing the resources of the operator. Each object represents a
        quantum gate and the number of times it occurs in the decomposition.

        Args:
            cdf_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.CDFHamiltonian`): a compressed double factorized
                Hamiltonian to be approximately exponentiated
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even.

        Resources:
            The resources are defined according to the recursive formula presented above.
            The number of times an operator, :math:`e^{itO_{j}}`, is applied depends on the
            number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

            .. math::

                C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

            Furthermore, because of the symmetric form of the recursive formula, the first and last terms get grouped.
            This reduces the counts for those terms to:

            .. math::

                \begin{align}
                    C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                    C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
                \end{align}

            The resources for a single step expansion of compressed double factorized Hamiltonian are
            calculated based on `arXiv:2506.15784 <https://arxiv.org/abs/2506.15784>`_.


        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.
        """
        k = order // 2
        gate_list = []
        num_orb = cdf_ham.num_orbitals
        num_frags = cdf_ham.num_fragments

        op_onebody = resource_rep(
            Prod,
            {"cmpr_factors_and_counts": ((RZ.resource_rep(), 2 * num_orb),)},
        )

        op_twobody = resource_rep(
            Prod,
            {
                "cmpr_factors_and_counts": (
                    (MultiRZ.resource_rep(num_wires=2), (2 * num_orb - 1) * num_orb),
                )
            },
        )

        basis_rot = resource_rep(BasisRotation, {"dim": num_orb})

        if order == 1:
            gate_list.append(GateCount(basis_rot, 2 * num_frags * num_steps))

            gate_list.append(GateCount(op_onebody, num_steps))
            gate_list.append(GateCount(op_twobody, (num_frags - 1) * num_steps))
            return gate_list

        # For first and last fragment
        gate_list.append(GateCount(basis_rot, 4 * num_steps * (5 ** (k - 1)) + 2))
        gate_list.append(GateCount(op_onebody, num_steps * (5 ** (k - 1)) + 1))
        gate_list.append(GateCount(op_twobody, num_steps * (5 ** (k - 1))))

        # For rest of the fragments
        gate_list.append(GateCount(basis_rot, 4 * num_steps * (num_frags - 2) * (5 ** (k - 1))))
        gate_list.append(GateCount(op_twobody, 2 * num_steps * (num_frags - 2) * (5 ** (k - 1))))

        return gate_list

    @classmethod
    def controlled_resource_decomp(
        cls, num_ctrl_wires: int, num_zero_ctrl: int, target_resource_params: dict | None = None
    ):
        r"""Returns the controlled resource decomposition.

        Args:
            num_ctrl_wires (int): the number of qubits the operation is controlled on
            num_zero_ctrl (int): the number of control qubits, that are controlled when in the :math:`|0\rangle` state
            target_resource_params (dict): dictionary containing the size of the larger of the two registers being added together

        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.

        Resources:
            The original resources are controlled only on the Z rotation gates.
        """
        cdf_ham = target_resource_params["cdf_ham"]
        num_steps = target_resource_params["num_steps"]
        order = target_resource_params["order"]

        k = order // 2
        gate_list = []
        num_orb = cdf_ham.num_orbitals
        num_frags = cdf_ham.num_fragments

        op_onebody = resource_rep(
            Prod,
            {
                "cmpr_factors_and_counts": [
                    (
                        resource_rep(
                            Controlled,
                            {
                                "base_cmpr_op": RZ.resource_rep(),
                                "num_ctrl_wires": num_ctrl_wires,
                                "num_zero_ctrl": num_zero_ctrl,
                            },
                        ),
                        (2 * num_orb),
                    )
                ]
            },
        )

        op_twobody = resource_rep(
            Prod,
            {
                "cmpr_factors_and_counts": [
                    (
                        resource_rep(
                            Controlled,
                            {
                                "base_cmpr_op": MultiRZ.resource_rep(num_wires=2),
                                "num_ctrl_wires": num_ctrl_wires,
                                "num_zero_ctrl": num_zero_ctrl,
                            },
                        ),
                        (2 * num_orb - 1) * num_orb,
                    )
                ]
            },
        )

        basis_rot = resource_rep(BasisRotation, {"dim": num_orb})

        if order == 1:
            gate_list.append(GateCount(basis_rot, 2 * num_frags * num_steps))

            gate_list.append(GateCount(op_onebody, num_steps))
            gate_list.append(GateCount(op_twobody, (num_frags - 1) * num_steps))
            return gate_list

        # For first and last fragment
        gate_list.append(GateCount(basis_rot, 4 * num_steps * (5 ** (k - 1)) + 2))
        gate_list.append(GateCount(op_onebody, num_steps * (5 ** (k - 1)) + 1))
        gate_list.append(GateCount(op_twobody, num_steps * (5 ** (k - 1))))

        # For rest of the fragments
        gate_list.append(GateCount(basis_rot, 4 * num_steps * (num_frags - 2) * (5 ** (k - 1))))
        gate_list.append(GateCount(op_twobody, 2 * num_steps * (num_frags - 2) * (5 ** (k - 1))))

        return gate_list


class TrotterTHC(ResourceOperator):
    r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
    exponential of a tensor hypercontracted (THC) Hamiltonian.

    The Suzuki-Trotter product formula provides a method to approximate the matrix exponential of
    Hamiltonian expressed as a linear combination of terms which in general do not commute.
    Consider the Hamiltonian :math:`H = \Sigma^{N}_{j=0} O_{j}`: the product formula is constructed using
    symmetrized products of the terms in the Hamiltonian. The symmetrized products of order
    :math:`m \in [1, 2, 4, ..., 2k]` with :math:`k \in \mathbb{N}` are given by:

    .. math::

        \begin{align}
            S_{1}(t) &= \Pi_{j=0}^{N} \ e^{i t O_{j}} \\
            S_{2}(t) &= \Pi_{j=0}^{N} \ e^{i \frac{t}{2} O_{j}} \cdot \Pi_{j=N}^{0} \ e^{i \frac{t}{2} O_{j}} \\
            &\vdots \\
            S_{m}(t) &= S_{m-2}(p_{m}t)^{2} \cdot S_{m-2}((1-4p_{m})t) \cdot S_{m-2}(p_{m}t)^{2},
        \end{align}

    where the coefficient is :math:`p_{m} = 1 / (4 - \sqrt[m - 1]{4})`. The :math:`m^{\text{th}}`
    order, :math:`n`-step Suzuki-Trotter approximation is then defined as:

    .. math::

        e^{iHt} \approx \left [S_{m}(t / n)  \right ]^{n}.

    For more details see `J. Math. Phys. 32, 400 (1991) <https://pubs.aip.org/aip/jmp/article-abstract/32/2/400/229229>`_.

    Args:
        thc_ham (:class:`~.pennylane.estimator.compact_hamiltonian.THCHamiltonian`): a tensor hypercontracted
            Hamiltonian to be approximately exponentiated
        num_steps (int): number of Trotter steps to perform
        order (int): order of the approximation, must be ``1`` or an even number
        wires (list[int] | None): the wires on which the operator acts

    Resources:
        The resources are defined according to the recursive formula presented above.
        The number of times an operator :math:`e^{itO_{j}}` is applied depends on the
        number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

        .. math::

            C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

        Furthermore, because of the symmetric form of the recursive formula, the first and last
        terms get grouped. This reduces the counts for those terms to:

        .. math::

            \begin{align}
                C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
            \end{align}

        The resources for a single step expansion of tensor hypercontracted Hamiltonian are
        calculated based on `arXiv:2407.04432 <https://arxiv.org/abs/2407.04432>`_.

    .. seealso::
        :class:`~.estimator.compact_hamiltonian.THCHamiltonian`

    .. seealso:: :class:`~.TrotterProduct`

    **Example**

    The resources for this operation are computed using:

    >>> import pennylane.estimator as qre
    >>> num_steps, order = (1, 2)
    >>> thc_ham = qre.THCHamiltonian(num_orbitals=4, tensor_rank=4)
    >>> res = qre.estimate(qre.TrotterTHC(thc_ham, num_steps, order))
    >>> print(res)
    --- Resources: ---
     Total wires: 8
        algorithmic wires: 8
        allocated wires: 0
             zero state: 0
             any state: 0
     Total gates : 8.520E+3
      'T': 7.888E+3,
      'CNOT': 128,
      'Z': 144,
      'S': 216,
      'Hadamard': 144
    """

    resource_keys = {"thc_ham", "num_steps", "order"}

    def __init__(
        self,
        thc_ham: THCHamiltonian,
        num_steps: int,
        order: int,
        wires: WiresLike | None = None,
    ):

        if not isinstance(thc_ham, THCHamiltonian):
            raise TypeError(
                f"Unsupported Hamiltonian representation for TrotterTHC."
                f"This method works with thc Hamiltonian, {type(thc_ham)} provided"
            )
        self.num_steps = num_steps
        self.order = order
        self.thc_ham = thc_ham

        self.num_wires = thc_ham.tensor_rank * 2

        if wires is not None and len(Wires(wires)) != self.num_wires:
            raise ValueError(f"Expected {self.num_wires} wires, got {len(Wires(wires))}")

        super().__init__(wires=wires)

    @property
    def resource_params(self) -> dict:
        r"""Returns a dictionary containing the minimal information needed to compute the resources.

        Returns:
            dict: A dictionary containing the resource parameters:
                * thc_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.THCHamiltonian`): a tensor hypercontracted
                  Hamiltonian to be approximately exponentiated
                * num_steps (int): number of Trotter steps to perform
                * order (int): order of the approximation, must be 1 or even
        """
        return {
            "thc_ham": self.thc_ham,
            "num_steps": self.num_steps,
            "order": self.order,
        }

    @classmethod
    def resource_rep(
        cls, thc_ham: THCHamiltonian, num_steps: int, order: int
    ) -> CompressedResourceOp:
        """Returns a compressed representation containing only the parameters of
        the Operator that are needed to compute the resources.

        Args:
            thc_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.THCHamiltonian`): a tensor hypercontracted
                Hamiltonian to be approximately exponentiated
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even

        Returns:
            :class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`: the operator in a compressed representation
        """
        params = {
            "thc_ham": thc_ham,
            "num_steps": num_steps,
            "order": order,
        }
        num_wires = thc_ham.tensor_rank * 2
        return CompressedResourceOp(cls, num_wires, params)

    @classmethod
    def resource_decomp(
        cls, thc_ham: THCHamiltonian, num_steps: int, order: int
    ) -> list[GateCount]:
        r"""Returns a list representing the resources of the operator. Each object represents a
        quantum gate and the number of times it occurs in the decomposition.

        Args:
            thc_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.THCHamiltonian`): a tensor hypercontracted
                Hamiltonian to be approximately exponentiated
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even

        Resources:
            The resources are defined according to the recursive formula presented above.
            The number of times an operator, :math:`e^{itO_{j}}`, is applied depends on the
            number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

            .. math::

                C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

            Furthermore, because of the symmetric form of the recursive formula, the first and last
            terms get grouped. This reduces the counts for those terms to:

            .. math::

                \begin{align}
                    C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                    C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
                \end{align}

            The resources for a single step expansion of tensor hypercontracted Hamiltonian are
            calculated based on `arXiv:2407.04432 <https://arxiv.org/abs/2407.04432>`_.


        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.
        """
        k = order // 2
        gate_list = []
        num_orb = thc_ham.num_orbitals
        tensor_rank = thc_ham.tensor_rank

        op_onebody = resource_rep(
            Prod,
            {"cmpr_factors_and_counts": ((RZ.resource_rep(), 2 * num_orb),)},
        )

        op_twobody = resource_rep(
            Prod,
            {
                "cmpr_factors_and_counts": (
                    (
                        MultiRZ.resource_rep(num_wires=2),
                        (2 * tensor_rank - 1) * tensor_rank,
                    ),
                )
            },
        )

        basis_rot_onebody = resource_rep(BasisRotation, {"dim": num_orb})
        basis_rot_twobody = resource_rep(BasisRotation, {"dim": tensor_rank})

        if order == 1:
            gate_list.append(GateCount(basis_rot_onebody, 2 * num_steps))
            gate_list.append(GateCount(basis_rot_twobody, 2 * num_steps))
            gate_list.append(GateCount(op_onebody, num_steps))
            gate_list.append(GateCount(op_twobody, num_steps))
            return gate_list

        # For one-body tensor
        gate_list.append(GateCount(basis_rot_onebody, 2 * num_steps * (5 ** (k - 1)) + 2))
        gate_list.append(GateCount(op_onebody, num_steps * (5 ** (k - 1)) + 1))

        # For two-body tensor
        gate_list.append(GateCount(basis_rot_twobody, 2 * num_steps * (5 ** (k - 1))))
        gate_list.append(GateCount(op_twobody, num_steps * (5 ** (k - 1))))

        return gate_list

    @classmethod
    def controlled_resource_decomp(
        cls, num_ctrl_wires: int, num_zero_ctrl: int, target_resource_params: dict | None = None
    ):
        r"""Returns the controlled resource decomposition.

        Args:
            num_ctrl_wires (int): the number of qubits the operation is controlled on
            num_zero_ctrl (int): the number of control qubits, that are controlled when in the :math:`|0\rangle` state
            target_resource_params (dict): dictionary containing the size of the larger of the two registers being added together

        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.

        Resources:
            The original resources are controlled only on the Z rotation gates
        """
        thc_ham = target_resource_params["thc_ham"]
        num_steps = target_resource_params["num_steps"]
        order = target_resource_params["order"]

        k = order // 2
        gate_list = []
        num_orb = thc_ham.num_orbitals
        tensor_rank = thc_ham.tensor_rank

        op_onebody = resource_rep(
            Prod,
            {
                "cmpr_factors_and_counts": [
                    (
                        resource_rep(
                            Controlled,
                            {
                                "base_cmpr_op": RZ.resource_rep(),
                                "num_ctrl_wires": num_ctrl_wires,
                                "num_zero_ctrl": num_zero_ctrl,
                            },
                        ),
                        (2 * num_orb),
                    )
                ]
            },
        )

        op_twobody = resource_rep(
            Prod,
            {
                "cmpr_factors_and_counts": [
                    (
                        resource_rep(
                            Controlled,
                            {
                                "base_cmpr_op": MultiRZ.resource_rep(num_wires=2),
                                "num_ctrl_wires": num_ctrl_wires,
                                "num_zero_ctrl": num_zero_ctrl,
                            },
                        ),
                        (2 * tensor_rank - 1) * tensor_rank,
                    )
                ]
            },
        )

        basis_rot_onebody = resource_rep(BasisRotation, {"dim": num_orb})
        basis_rot_twobody = resource_rep(BasisRotation, {"dim": tensor_rank})

        if order == 1:
            gate_list.append(GateCount(basis_rot_onebody, 2 * num_steps))
            gate_list.append(GateCount(basis_rot_twobody, 2 * num_steps))
            gate_list.append(GateCount(op_onebody, num_steps))
            gate_list.append(GateCount(op_twobody, num_steps))
            return gate_list

        # For one-body tensor
        gate_list.append(GateCount(basis_rot_onebody, 2 * num_steps * (5 ** (k - 1)) + 2))
        gate_list.append(GateCount(op_onebody, num_steps * (5 ** (k - 1)) + 1))

        # For two-body tensor
        gate_list.append(GateCount(basis_rot_twobody, 2 * num_steps * (5 ** (k - 1))))
        gate_list.append(GateCount(op_twobody, num_steps * (5 ** (k - 1))))

        return gate_list


class TrotterVibrational(ResourceOperator):
    r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
    exponential of a vibrational Hamiltonian.

    The Suzuki-Trotter product formula provides a method to approximate the matrix exponential of
    Hamiltonian expressed as a linear combination of terms which in general do not commute.
    Consider the Hamiltonian :math:`H = \Sigma^{N}_{j=0} O_{j}`: the product formula is constructed using
    symmetrized products of the terms in the Hamiltonian. The symmetrized products of order
    :math:`m \in [1, 2, 4, ..., 2k]` with :math:`k \in \mathbb{N}` are given by:

    .. math::

        \begin{align}
            S_{1}(t) &= \Pi_{j=0}^{N} \ e^{i t O_{j}} \\
            S_{2}(t) &= \Pi_{j=0}^{N} \ e^{i \frac{t}{2} O_{j}} \cdot \Pi_{j=N}^{0} \ e^{i \frac{t}{2} O_{j}} \\
            &\vdots \\
            S_{m}(t) &= S_{m-2}(p_{m}t)^{2} \cdot S_{m-2}((1-4p_{m})t) \cdot S_{m-2}(p_{m}t)^{2},
        \end{align}

    where the coefficient is :math:`p_{m} = 1 / (4 - \sqrt[m - 1]{4})`. The :math:`m^\text{th}`
    order, :math:`n`-step Suzuki-Trotter approximation is then defined as:

    .. math::

        e^{iHt} \approx \left [S_{m}(t / n)  \right ]^{n}.

    For more details see `J. Math. Phys. 32, 400 (1991) <https://pubs.aip.org/aip/jmp/article-abstract/32/2/400/229229>`_.

    Args:
        vibration_ham (:class:`~.pennylane.estimator.compact_hamiltonian.VibrationalHamiltonian`): a real space vibrational
            Hamiltonian to be approximately exponentiated
        num_steps (int): number of Trotter steps to perform
        order (int): order of the approximation, must be ``1`` or an even number
        phase_grad_precision (float | None): precision for the phase gradient calculation
        coeff_precision (float | None): precision for the loading of coefficients
        wires (list[int] | None): the wires on which the operator acts

    Resources:
        The resources are defined according to the recursive formula presented above.
        The number of times an operator :math:`e^{itO_{j}}` is applied depends on the
        number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

        .. math::

            C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

        Furthermore, because of the symmetric form of the recursive formula, the first and last terms get grouped.
        This reduces the counts for those terms to:

        .. math::

            \begin{align}
                C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
            \end{align}

        The resources for a single step expansion of vibrational Hamiltonian are calculated based on
        `arXiv:2504.10602 <https://arxiv.org/pdf/2504.10602>`_.

    .. seealso::
        :class:`~.estimator.compact_hamiltonian.VibrationalHamiltonian`

    .. seealso:: :class:`~.TrotterProduct`

    **Example**

    The resources for this operation are computed using:

    >>> import pennylane.estimator as qre
    >>> num_steps, order = (10, 2)
    >>> vibration_ham = qre.VibrationalHamiltonian(num_modes=2, grid_size=4, taylor_degree=2)
    >>> res = qre.estimate(qre.TrotterVibrational(vibration_ham, num_steps, order))
    >>> print(res)
    --- Resources: ---
     Total wires: 83
        algorithmic wires: 8
        allocated wires: 75
             zero state: 75
             any state: 0
     Total gates : 1.239E+5
      'Toffoli': 2.248E+4,
      'T': 749,
      'CNOT': 3.520E+4,
      'X': 1.216E+3,
      'Z': 1,
      'S': 1,
      'Hadamard': 6.422E+4
    """

    resource_keys = {
        "vibration_ham",
        "num_steps",
        "order",
        "phase_grad_precision",
        "coeff_precision",
    }

    def __init__(
        self,
        vibration_ham: VibrationalHamiltonian,
        num_steps: int,
        order: int,
        phase_grad_precision: float | None = None,
        coeff_precision: float | None = None,
        wires: WiresLike | None = None,
    ):

        if not isinstance(vibration_ham, VibrationalHamiltonian):
            raise TypeError(
                f"Unsupported Hamiltonian representation for TrotterVibrational."
                f"This method works with vibrational Hamiltonian, {type(vibration_ham)} provided"
            )

        self.num_steps = num_steps
        self.order = order
        self.vibration_ham = vibration_ham
        self.phase_grad_precision = phase_grad_precision
        self.coeff_precision = coeff_precision

        self.num_wires = vibration_ham.num_modes * vibration_ham.grid_size

        if wires is not None and len(Wires(wires)) != self.num_wires:
            raise ValueError(f"Expected {self.num_wires} wires, got {len(Wires(wires))}")

        super().__init__(wires=wires)

    @property
    def resource_params(self) -> dict:
        r"""Returns a dictionary containing the minimal information needed to compute the resources.

        Returns:
            dict: A dictionary containing the resource parameters:
                * vibration_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.VibrationalHamiltonian`): a real space vibrational
                  Hamiltonian to be approximately exponentiated.
                * num_steps (int): number of Trotter steps to perform
                * order (int): order of the approximation, must be 1 or even
                * phase_grad_precision (float): precision for the phase gradient calculation,
                * coeff_precision (float): precision for the loading of coefficients,
        """
        return {
            "vibration_ham": self.vibration_ham,
            "num_steps": self.num_steps,
            "order": self.order,
            "phase_grad_precision": self.phase_grad_precision,
            "coeff_precision": self.coeff_precision,
        }

    @classmethod
    def resource_rep(
        cls,
        vibration_ham: VibrationalHamiltonian,
        num_steps: int,
        order: int,
        phase_grad_precision: float | None = None,
        coeff_precision: float | None = None,
    ) -> CompressedResourceOp:
        """Returns a compressed representation containing only the parameters of
        the Operator that are needed to compute the resources.

        Args:
            vibration_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.VibrationalHamiltonian`): a real space vibrational
                Hamiltonian to be approximately exponentiated.
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even
            phase_grad_precision (float | None): precision for the phase gradient calculation
            coeff_precision (float | None): precision for the loading of coefficients

        Returns:
            :class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`: the operator in a compressed representation
        """
        params = {
            "vibration_ham": vibration_ham,
            "num_steps": num_steps,
            "order": order,
            "phase_grad_precision": phase_grad_precision,
            "coeff_precision": coeff_precision,
        }
        num_wires = vibration_ham.num_modes * vibration_ham.grid_size
        return CompressedResourceOp(cls, num_wires, params)

    @staticmethod
    def _cached_terms(grid_size, taylor_degree, coeff_precision, cached_tree, path, index):
        r"""Recursive function to compute the resources for the trotterization of vibrational Hamiltonian
        while caching the coefficients."""

        cur_path, len_path = tuple(path), len(path)
        coeff_wires = int(abs(np.floor(np.log2(coeff_precision))))
        gate_cache = []

        x = X.resource_rep()
        if 1 < len_path <= taylor_degree and cur_path not in cached_tree[len_path]:

            if len(cached_tree[len_path]):
                prev_state = cached_tree[len_path][-1]

                if len_path == 2 and prev_state[0] == prev_state[1]:
                    out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size)
                    gate_cache.append(GateCount(out_square, 1))
                elif len_path == 4 and len(set(prev_state)) == 1:
                    out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size * 2)
                    gate_cache.append(GateCount(out_square, 1))
                else:
                    multiplier = OutMultiplier.resource_rep(grid_size, grid_size * (len_path - 1))
                    gate_cache.append(GateCount(multiplier, 1))

            # Add the Square / Multiplier for current state
            if len_path == 2 and cur_path[-1] == cur_path[-2]:
                out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size)
                gate_cache.append(GateCount(out_square, 1))
            elif len_path == 4 and len(set(cur_path)) == 1:
                out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size * 2)
                gate_cache.append(GateCount(out_square, 1))
            else:
                multiplier = OutMultiplier.resource_rep(grid_size, grid_size * (len_path - 1))
                gate_cache.append(GateCount(multiplier, 1))

            # Add the coefficient Initializer for current state
            # assuming that half the bits in the coefficient are 1
            gate_cache.append(GateCount(x, int(coeff_wires / 2)))

            # Add the Multiplier for current coefficient
            multiplier = OutMultiplier.resource_rep(grid_size * len_path, coeff_wires)
            gate_cache.append(GateCount(multiplier, 1))

            # Add the Adder for Resource state
            adder = SemiAdder.resource_rep(max_register_size=2 * max(coeff_wires, 2 * grid_size))
            gate_cache.append(GateCount(adder, 1))

            # Adjoint the Multiplier for current coefficient
            multiplier = OutMultiplier.resource_rep(grid_size * len_path, coeff_wires)
            gate_cache.append(GateCount(multiplier, 1))

            # Adjoint the coefficient Initializer for current state
            # assuming that half the bits in the coefficient are 1
            gate_cache.append(GateCount(x, coeff_wires / 2))

            cached_tree[len_path].append(cur_path)

        if len_path < taylor_degree and index + 1:
            gate_cache_curr, cached_tree = TrotterVibrational._cached_terms(
                grid_size, taylor_degree, coeff_precision, cached_tree, path + [index], index
            )  # Depth first search traversal with current element
            gate_cache += gate_cache_curr
            gate_cache_next, cached_tree = TrotterVibrational._cached_terms(
                grid_size, taylor_degree, coeff_precision, cached_tree, path, index - 1
            )  # Depth first search traversal with next element
            gate_cache += gate_cache_next

        return gate_cache, cached_tree

    @staticmethod
    def _rep_circuit(vibration_ham: VibrationalHamiltonian, coeff_precision, num_rep):
        r"""Returns the expansion of the circuit with given number of repetitions."""

        num_modes = vibration_ham.num_modes
        grid_size = vibration_ham.grid_size
        taylor_degree = vibration_ham.taylor_degree

        gate_lst = []
        # Shifted QFT for kinetic part

        t = T.resource_rep()
        gate_lst.append(GateCount(t, num_rep * (num_modes * int(np.ceil(np.log2(num_modes) - 1)))))

        kinetic_deg = 2
        cached_tree = {index: [] for index in range(1, kinetic_deg + 1)}
        gate_cache, cached_tree = TrotterVibrational._cached_terms(
            grid_size, kinetic_deg, coeff_precision, cached_tree, path=[], index=num_modes - 1
        )
        gate_lst += gate_cache * num_rep

        cached_tree = {index: [] for index in range(1, taylor_degree + 1)}
        gate_cache, cached_tree = TrotterVibrational._cached_terms(
            grid_size, taylor_degree, coeff_precision, cached_tree, path=[], index=num_modes - 1
        )
        gate_lst += gate_cache * num_rep

        # Adjoints for the last Squares / Multipliers
        for idx in range(2, taylor_degree):
            last_state = cached_tree[idx][-1]
            if idx == 2 and last_state[-1] == last_state[-2]:
                gate_lst.append(
                    GateCount(OutOfPlaceSquare.resource_rep(register_size=grid_size), num_rep)
                )
            elif idx == 4 and len(set(last_state)) == 1:
                gate_lst.append(
                    GateCount(
                        OutOfPlaceSquare.resource_rep(register_size=grid_size * 2),
                        num_rep,
                    )
                )
            else:
                gate_lst.append(
                    GateCount(
                        OutMultiplier.resource_rep(grid_size, grid_size * (idx - 1)),
                        num_rep,
                    )
                )

        # Shifted QFT Adjoint
        gate_lst.append(GateCount(t, num_rep * (num_modes * int(np.ceil(np.log2(num_modes) - 1)))))

        return gate_lst

    @classmethod
    def resource_decomp(
        cls,
        vibration_ham: VibrationalHamiltonian,
        num_steps: int,
        order: int,
        phase_grad_precision: float | None = None,
        coeff_precision: float | None = None,
    ) -> list[GateCount]:
        r"""Returns a list representing the resources of the operator. Each object represents a quantum gate
        and the number of times it occurs in the decomposition.

        Args:
            vibration_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.VibrationalHamiltonian`): a real space vibrational
                Hamiltonian to be approximately exponentiated.
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even
            phase_grad_precision (float | None): precision for the phase gradient calculation
            coeff_precision (float | None): precision for the loading of coefficients

        Resources:
            The resources are defined according to the recursive formula presented above.
            The number of times an operator, :math:`e^{itO_{j}}`, is applied depends on the
            number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

            .. math::

                C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

            Furthermore, because of the symmetric form of the recursive formula, the first and last terms get grouped.
            This reduces the counts for those terms to:

            .. math::

                \begin{align}
                    C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                    C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
                \end{align}

            The resources for a single step expansion of vibrational Hamiltonian are calculated based on
            `arXiv:2504.10602 <https://arxiv.org/pdf/2504.10602>`_.

        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.
        """

        k = order // 2
        gate_list = []
        num_modes = vibration_ham.num_modes
        grid_size = vibration_ham.grid_size
        taylor_degree = vibration_ham.taylor_degree

        phase_grad_wires = int(abs(np.floor(np.log2(phase_grad_precision))))
        coeff_wires = int(abs(np.floor(np.log2(coeff_precision))))

        x = X.resource_rep()

        phase_grad = PhaseGradient.resource_rep(phase_grad_wires)

        # Allocate the phase gradient registers
        gate_list.append(Allocate(phase_grad_wires * (taylor_degree - 1)))
        # Resource Registers
        gate_list.append(GateCount(phase_grad, taylor_degree - 1))

        # Allocate auxiliary registers for the coefficients
        gate_list.append(Allocate(4 * grid_size + 2 * coeff_wires))

        # Basis state prep per mode, implemented only for the first step
        gate_list.append(GateCount(x, num_modes * grid_size))

        if order == 1:
            gate_list += TrotterVibrational._rep_circuit(vibration_ham, coeff_precision, num_steps)
        else:
            gate_list += TrotterVibrational._rep_circuit(
                vibration_ham, coeff_precision, 2 * num_steps * (5 ** (k - 1))
            )

        # Adjoint of Basis state prep, implemented only for the last step
        gate_list.append(GateCount(x, num_modes * grid_size))

        # Free auxiliary registers for the coefficients
        gate_list.append(Deallocate(4 * grid_size + 2 * coeff_wires))

        # Deallocate the phase gradient registers
        gate_list.append(Deallocate(phase_grad_wires * (taylor_degree - 1)))

        return gate_list


class TrotterVibronic(ResourceOperator):
    r"""An operation representing the Suzuki-Trotter product approximation for the complex matrix
    exponential of a real-space vibronic Hamiltonian.

    The Suzuki-Trotter product formula provides a method to approximate the matrix exponential of
    Hamiltonian expressed as a linear combination of terms which in general do not commute.
    Consider the Hamiltonian :math:`H = \Sigma^{N}_{j=0} O_{j}`: the product formula is constructed using
    symmetrized products of the terms in the Hamiltonian. The symmetrized products of order
    :math:`m \in [1, 2, 4, ..., 2k]` with :math:`k \in \mathbb{N}` are given by:

    .. math::

        \begin{align}
            S_{1}(t) &= \Pi_{j=0}^{N} \ e^{i t O_{j}} \\
            S_{2}(t) &= \Pi_{j=0}^{N} \ e^{i \frac{t}{2} O_{j}} \cdot \Pi_{j=N}^{0} \ e^{i \frac{t}{2} O_{j}} \\
            &\vdots \\
            S_{m}(t) &= S_{m-2}(p_{m}t)^{2} \cdot S_{m-2}((1-4p_{m})t) \cdot S_{m-2}(p_{m}t)^{2},
        \end{align}

    where the coefficient is :math:`p_{m} = 1 / (4 - \sqrt[m - 1]{4})`. The :math:`m^{\text{th}}`
    order, :math:`n`-step Suzuki-Trotter approximation is then defined as:

    .. math::

        e^{iHt} \approx \left [S_{m}(t / n)  \right ]^{n}.

    For more details see `J. Math. Phys. 32, 400 (1991) <https://pubs.aip.org/aip/jmp/article-abstract/32/2/400/229229>`_.

    Args:
        vibronic_ham (:class:`~.pennylane.estimator.compact_hamiltonian.VibronicHamiltonian`): a real-space vibronic
            Hamiltonian to be approximately exponentiated
        num_steps (int): number of Trotter steps to perform
        order (int): order of the approximation, must be ``1`` or an even number
        phase_grad_precision (float | None): precision for the phase gradient calculation
        coeff_precision (float | None): precision for the loading of coefficients
        wires (list[int] | None): the wires on which the operator acts.

    Resources:
        The resources are defined according to the recursive formula presented above.
        The number of times an operator :math:`e^{itO_{j}}` is applied depends on the
        number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

        .. math::

            C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

        Furthermore, because of the symmetric form of the recursive formula, the first and last terms get grouped.
        This reduces the counts for those terms to:

        .. math::

            \begin{align}
                C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
            \end{align}

        The resources for a single step expansion of real-space vibronic Hamiltonian are calculated
        based on `arXiv:2411.13669 <https://arxiv.org/abs/2411.13669>`_.

    .. seealso::
        :class:`~.estimator.compact_hamiltonian.VibronicHamiltonian`

    .. seealso:: :class:`~.TrotterProduct`

    **Example**

    The resources for this operation are computed using:

    >>> import pennylane.estimator as qre
    >>> num_steps, order = (10, 2)
    >>> vibronic_ham = qre.VibronicHamiltonian(num_modes=2, num_states=4, grid_size=4, taylor_degree=2)
    >>> res = qre.estimate(qre.TrotterVibronic(vibronic_ham, num_steps, order))
    >>> print(res)
    --- Resources: ---
     Total wires: 85
        algorithmic wires: 10
        allocated wires: 75
             zero state: 75
             any state: 0
     Total gates : 1.332E+5
      'Toffoli': 2.320E+4,
      'T': 749,
      'CNOT': 4.144E+4,
      'X': 1.456E+3,
      'Z': 1,
      'S': 1,
      'Hadamard': 6.638E+4
    """

    resource_keys = {
        "vibronic_ham",
        "num_steps",
        "order",
        "phase_grad_precision",
        "coeff_precision",
    }

    def __init__(
        self,
        vibronic_ham: VibronicHamiltonian,
        num_steps: int,
        order: int,
        phase_grad_precision: float | None = None,
        coeff_precision: float | None = None,
        wires: WiresLike | None = None,
    ):

        if not isinstance(vibronic_ham, VibronicHamiltonian):
            raise TypeError(
                f"Unsupported Hamiltonian representation for TrotterVibronic."
                f"This method works with vibronic Hamiltonian, {type(vibronic_ham)} provided"
            )

        self.num_steps = num_steps
        self.order = order
        self.vibronic_ham = vibronic_ham
        self.phase_grad_precision = phase_grad_precision
        self.coeff_precision = coeff_precision

        self.num_wires = (
            int(np.ceil(np.log2(vibronic_ham.num_states)))
            + vibronic_ham.num_modes * vibronic_ham.grid_size
        )

        if wires is not None and len(Wires(wires)) != self.num_wires:
            raise ValueError(f"Expected {self.num_wires} wires, got {len(Wires(wires))}")

        super().__init__(wires=wires)

    @property
    def resource_params(self) -> dict:
        r"""Returns a dictionary containing the minimal information needed to compute the resources.

        Returns:
            dict: A dictionary containing the resource parameters:
                * vibronic_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.VibronicHamiltonian`): a real-space vibronic
                  Hamiltonian to be approximately exponentiated
                * num_steps (int): number of Trotter steps to perform
                * order (int): order of the approximation, must be 1 or even
                * phase_grad_precision (float): precision for the phase gradient calculation
                * coeff_precision (float): precision for the loading of coefficients
        """
        return {
            "vibronic_ham": self.vibronic_ham,
            "num_steps": self.num_steps,
            "order": self.order,
            "phase_grad_precision": self.phase_grad_precision,
            "coeff_precision": self.coeff_precision,
        }

    @classmethod
    def resource_rep(
        cls,
        vibronic_ham: VibronicHamiltonian,
        num_steps: int,
        order: int,
        phase_grad_precision: float | None = None,
        coeff_precision: float | None = None,
    ) -> CompressedResourceOp:
        """Returns a compressed representation containing only the parameters of
        the Operator that are needed to compute a resource estimation.

        Args:
            vibronic_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.VibronicHamiltonian`): a real space vibronic
                Hamiltonian to be approximately exponentiated
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even
            phase_grad_precision (float | None): precision for the phase gradient calculation
            coeff_precision (float | None): precision for the loading of coefficients
        Returns:
            :class:`~.pennylane.estimator.resource_operator.CompressedResourceOp`: the operator in a compressed representation
        """
        params = {
            "vibronic_ham": vibronic_ham,
            "num_steps": num_steps,
            "order": order,
            "phase_grad_precision": phase_grad_precision,
            "coeff_precision": coeff_precision,
        }
        num_wires = (
            int(np.ceil(np.log2(vibronic_ham.num_states)))
            + vibronic_ham.num_modes * vibronic_ham.grid_size
        )
        return CompressedResourceOp(cls, num_wires, params)

    @staticmethod
    def _cached_terms(
        num_states, grid_size, taylor_degree, coeff_precision, cached_tree, path, index
    ):
        r"""Recursive function to compute the resources for the trotterization of vibronic Hamiltonian
        while caching the coefficients."""

        cur_path, len_path = tuple(path), len(path)
        coeff_wires = int(abs(int(np.floor(np.log2(coeff_precision)))))
        gate_cache = []

        if 1 < len_path <= taylor_degree and cur_path not in cached_tree[len_path]:

            if len(cached_tree[len_path]):
                prev_state = cached_tree[len_path][-1]

                if len_path == 2 and prev_state[0] == prev_state[1]:
                    out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size)
                    gate_cache.append(GateCount(out_square, 1))
                elif len_path == 4 and len(set(prev_state)) == 1:
                    out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size * 2)
                    gate_cache.append(GateCount(out_square, 1))
                else:
                    multiplier = OutMultiplier.resource_rep(grid_size, grid_size * (len_path - 1))
                    gate_cache.append(GateCount(multiplier, 1))

            # Add the Square / Multiplier for current state
            if len_path == 2 and cur_path[-1] == cur_path[-2]:
                out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size)
                gate_cache.append(GateCount(out_square, 1))
            elif len_path == 4 and len(set(cur_path)) == 1:
                out_square = OutOfPlaceSquare.resource_rep(register_size=grid_size * 2)
                gate_cache.append(GateCount(out_square, 1))
            else:
                multiplier = OutMultiplier.resource_rep(grid_size, grid_size * (len_path - 1))
                gate_cache.append(GateCount(multiplier, 1))

            # Add the coefficient Initializer for current state
            # assuming that half the bits in the coefficient are 1
            coeff_unitaries = (
                resource_rep(
                    Prod,
                    {"cmpr_factors_and_counts": ((X.resource_rep(), int(coeff_wires / 2)),)},
                ),
            ) * num_states

            select_op = resource_rep(Select, {"cmpr_ops": coeff_unitaries})
            gate_cache.append(GateCount(select_op, 1))

            # Add the Multiplier for current coefficient
            multiplier = OutMultiplier.resource_rep(grid_size * len_path, coeff_wires)
            gate_cache.append(GateCount(multiplier, 1))

            # Add the Adder for Resource state
            adder = SemiAdder.resource_rep(max_register_size=2 * max(coeff_wires, 2 * grid_size))
            gate_cache.append(GateCount(adder, 1))

            # Adjoint the Multiplier for current coefficient
            multiplier = OutMultiplier.resource_rep(grid_size * len_path, coeff_wires)
            gate_cache.append(GateCount(multiplier, 1))

            # Adjoint the coefficient Initializer for current state
            # assuming that half the bits in the coefficient are 1
            gate_cache.append(GateCount(select_op, 1))

            cached_tree[len_path].append(cur_path)

        if len_path < taylor_degree and index + 1:
            gate_cache_curr, cached_tree = TrotterVibronic._cached_terms(
                num_states,
                grid_size,
                taylor_degree,
                coeff_precision,
                cached_tree,
                path + [index],
                index,
            )  # DFS with current element
            gate_cache += gate_cache_curr
            gate_cache_next, cached_tree = TrotterVibronic._cached_terms(
                num_states, grid_size, taylor_degree, coeff_precision, cached_tree, path, index - 1
            )  # DFS with next element
            gate_cache += gate_cache_next

        return gate_cache, cached_tree

    @staticmethod
    def _rep_circuit(vibronic_ham: VibronicHamiltonian, coeff_precision, num_rep):
        r"""Returns the expansion of the circuit with given number of repetitions."""

        num_modes = vibronic_ham.num_modes
        num_states = vibronic_ham.num_states
        grid_size = vibronic_ham.grid_size
        taylor_degree = vibronic_ham.taylor_degree

        gate_lst = []
        # Shifted QFT for kinetic part
        t = T.resource_rep()
        gate_lst.append(GateCount(t, num_rep * (num_modes * int(np.ceil(np.log2(num_modes) - 1)))))

        kinetic_deg = 2
        cached_tree = {index: [] for index in range(1, kinetic_deg + 1)}
        gate_cache, cached_tree = TrotterVibronic._cached_terms(
            num_states,
            grid_size,
            kinetic_deg,
            coeff_precision,
            cached_tree,
            path=[],
            index=num_modes - 1,
        )
        gate_lst += gate_cache * num_rep

        cached_tree = {index: [] for index in range(1, taylor_degree + 1)}
        gate_cache, cached_tree = TrotterVibronic._cached_terms(
            num_states,
            grid_size,
            taylor_degree,
            coeff_precision,
            cached_tree,
            path=[],
            index=num_modes - 1,
        )
        gate_lst += gate_cache * num_rep

        # Adjoints for the last Squares / Multipliers
        for idx in range(2, taylor_degree):
            last_state = cached_tree[idx][-1]
            if idx == 2 and last_state[-1] == last_state[-2]:
                gate_lst.append(
                    GateCount(OutOfPlaceSquare.resource_rep(register_size=grid_size), num_rep)
                )
            elif idx == 4 and len(set(last_state)) == 1:
                gate_lst.append(
                    GateCount(
                        OutOfPlaceSquare.resource_rep(register_size=grid_size * 2),
                        num_rep,
                    )
                )
            else:
                gate_lst.append(
                    GateCount(
                        OutMultiplier.resource_rep(grid_size, grid_size * (idx - 1)),
                        num_rep,
                    )
                )

        # Shifted QFT Adjoint
        gate_lst.append(GateCount(t, num_rep * (num_modes * int(np.ceil(np.log2(num_modes) - 1)))))

        return gate_lst

    @classmethod
    def resource_decomp(
        cls,
        vibronic_ham: VibronicHamiltonian,
        num_steps: int,
        order: int,
        phase_grad_precision: float | None,
        coeff_precision: float | None,
    ) -> list[GateCount]:
        r"""Returns a list representing the resources of the operator. Each object represents a quantum gate
        and the number of times it occurs in the decomposition.

        Args:
            vibronic_ham (:class:`~.pennylane.estimator.templates.compact_hamiltonian.VibronicHamiltonian`): a real space vibronic
                Hamiltonian to be approximately exponentiated
            num_steps (int): number of Trotter steps to perform
            order (int): order of the approximation, must be 1 or even
            phase_grad_precision (float | None): precision for the phase gradient calculation
            coeff_precision (float | None): precision for the loading of coefficients

        Resources:
            The resources are defined according to the recursive formula presented above.
            The number of times an operator, :math:`e^{itO_{j}}`, is applied depends on the
            number of Trotter steps (`n`) and the order of the approximation (`m`) and is given by:

            .. math::

                C_{O_j} = 2 * n \cdot 5^{\frac{m}{2} - 1}.

            Furthermore, because of the symmetric form of the recursive formula, the first and last terms get grouped.
            This reduces the counts for those terms to:

            .. math::

                \begin{align}
                    C_{O_{0}} &= n \cdot 5^{\frac{m}{2} - 1} + 1,  \\
                    C_{O_{N}} &= n \cdot 5^{\frac{m}{2} - 1}.
                \end{align}

            The resources for a single step expansion of real-space vibronic Hamiltonian are calculated
            based on `arXiv:2411.13669 <https://arxiv.org/abs/2411.13669>`_.

        Returns:
            list[:class:`~.pennylane.estimator.resource_operator.GateCount`]: A list of GateCount objects, where each object
            represents a specific quantum gate and the number of times it appears
            in the decomposition.
        """

        k = order // 2
        gate_list = []
        num_modes = vibronic_ham.num_modes
        num_states = vibronic_ham.num_states
        grid_size = vibronic_ham.grid_size
        taylor_degree = vibronic_ham.taylor_degree

        phase_grad_wires = int(abs(np.floor(np.log2(phase_grad_precision))))
        coeff_wires = int(abs(np.floor(np.log2(coeff_precision))))

        x = X.resource_rep()

        phase_grad = PhaseGradient.resource_rep(phase_grad_wires)

        # Allocate the phase gradient registers
        gate_list.append(Allocate(phase_grad_wires * (taylor_degree - 1)))
        # Resource Registers
        gate_list.append(GateCount(phase_grad, taylor_degree - 1))

        # Allocate auxiliary registers for the coefficients
        gate_list.append(Allocate(4 * grid_size + 2 * coeff_wires))

        # Basis state prep per mode, implemented only for the first step
        gate_list.append(GateCount(x, num_modes * grid_size))

        # electronic state
        gate_list.append(GateCount(resource_rep(Hadamard), int(np.ceil(np.log2(num_states)))))

        if order == 1:
            gate_list += TrotterVibronic._rep_circuit(vibronic_ham, coeff_precision, num_steps)
        else:
            gate_list += TrotterVibronic._rep_circuit(
                vibronic_ham, coeff_precision, 2 * num_steps * (5 ** (k - 1))
            )

        # Adjoint for electronic state
        gate_list.append(GateCount(resource_rep(Hadamard), int(np.ceil(np.log2(num_states)))))

        # Adjoint of Basis state prep, implemented only for the first step
        gate_list.append(GateCount(x, num_modes * grid_size))

        # Free auxiliary registers for the coefficients
        gate_list.append(Deallocate(4 * grid_size + 2 * coeff_wires))

        # Deallocate the phase gradient registers
        gate_list.append(Deallocate(phase_grad_wires * (taylor_degree - 1)))

        return gate_list
